package limit

import (
	"runtime"
	"sync"
	"sync/atomic"
)

type Limiter interface {
	ApplyForResources(uint32) bool //申请大小
	ReleaseResource(uint32)        //释放
	Close()
}
type bandwidthlimit struct {
	maxLimit uint32 //最大值
	applyNum uint32 //已申请的值
	mu       sync.Mutex
	isClose  int32
}

const (
	defaultMaxBytes = 1 << 26 //默认最大下载带宽64M
)

/**
 * @description: 申请资源
 * @param {uint32} byteNum
 * @return {*}
 */
func (bw *bandwidthlimit) ApplyForResources(byteNum uint32) bool {
	if atomic.LoadInt32(&bw.isClose) > 0 { //已关闭默认返回true
		return true
	}
	bw.mu.Lock()
	for {
		if byteNum > bw.maxLimit { //如果单个文件大小大于最大值
			if atomic.CompareAndSwapUint32(&bw.applyNum, 0, byteNum) {
				bw.mu.Unlock()
				return true
			}
			runtime.Gosched()
		}
		if atomic.LoadUint32(&bw.applyNum)+byteNum > bw.maxLimit {
			runtime.Gosched() //让出cpu等待下次执行
			continue
		}
		atomic.AddUint32(&bw.applyNum, byteNum)
		bw.mu.Unlock()
		return true
	}
}

/**
 * @description: 释放资源
 * @param {uint32} byteNum
 * @return {*}
 */
func (bw *bandwidthlimit) ReleaseResource(byteNum uint32) {
	if atomic.LoadInt32(&bw.isClose) > 0 {
		return
	}
	atomic.AddUint32(&bw.applyNum, -byteNum)
}

func (bw *bandwidthlimit) Close() {
	atomic.CompareAndSwapInt32(&bw.isClose, 0, 1)
}

/**
 * @description: 限制下载带宽，传0默认最大64M
 * @param {uint32} maxLimit
 * @return {*}
 */
func NewBandwidthlimiter(maxLimit uint32) Limiter {
	if maxLimit == 0 {
		maxLimit = defaultMaxBytes
	}
	return &bandwidthlimit{maxLimit: maxLimit}
}
